{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 001 | Val Acc: 55.58%\n",
      "Epoch 002 | Val Acc: 61.83%\n",
      "Epoch 003 | Val Acc: 60.00%\n",
      "Epoch 004 | Val Acc: 60.08%\n",
      "Epoch 005 | Val Acc: 59.08%\n",
      "Epoch 006 | Val Acc: 60.42%\n",
      "Epoch 007 | Val Acc: 58.92%\n",
      "Epoch 008 | Val Acc: 58.00%\n",
      "Epoch 009 | Val Acc: 64.17%\n",
      "Epoch 010 | Val Acc: 65.08%\n",
      "Epoch 011 | Val Acc: 60.25%\n",
      "Epoch 012 | Val Acc: 64.67%\n",
      "Epoch 013 | Val Acc: 58.08%\n",
      "Epoch 014 | Val Acc: 65.92%\n",
      "Epoch 015 | Val Acc: 60.08%\n",
      "Epoch 016 | Val Acc: 60.67%\n",
      "Epoch 017 | Val Acc: 58.25%\n",
      "Epoch 018 | Val Acc: 59.17%\n",
      "Epoch 019 | Val Acc: 59.92%\n",
      "Epoch 020 | Val Acc: 66.92%\n",
      "Epoch 021 | Val Acc: 51.25%\n",
      "Epoch 022 | Val Acc: 65.50%\n",
      "Epoch 023 | Val Acc: 62.92%\n",
      "Epoch 024 | Val Acc: 50.75%\n",
      "Epoch 025 | Val Acc: 62.42%\n",
      "Epoch 026 | Val Acc: 59.33%\n",
      "Epoch 027 | Val Acc: 62.08%\n",
      "Epoch 028 | Val Acc: 56.17%\n",
      "Epoch 029 | Val Acc: 58.25%\n",
      "Epoch 030 | Val Acc: 62.08%\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "import numpy as np\n",
    "\n",
    "# 加载数据\n",
    "flux_data = np.load('E:\\\\workplace2\\\\place\\\\spec_cnn\\\\flux.npy')  # (6000, 7781)\n",
    "spectypes = np.load('E:\\\\workplace2\\\\place\\\\spec_cnn\\\\spectypes.npy')  # (6000,)\n",
    "\n",
    "# 转换为 PyTorch 张量并调整形状\n",
    "train = torch.tensor(flux_data, dtype=torch.float32)  # (6000, 7781)\n",
    "target = torch.tensor(spectypes, dtype=torch.long)    # (6000,)\n",
    "\n",
    "# 调整为 RNN 输入格式 (batch, seq_len, input_dim)\n",
    "train = train.view(-1, 1, 7781)  # (6000, 1, 7781)\n",
    "\n",
    "# 自定义数据集\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data, labels):\n",
    "        self.data = data\n",
    "        self.labels = labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return self.data[idx], self.labels[idx]\n",
    "\n",
    "# 创建数据集和加载器\n",
    "full_dataset = CustomDataset(train, target)\n",
    "train_size = int(0.8 * len(full_dataset))\n",
    "test_size = len(full_dataset) - train_size\n",
    "train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size])\n",
    "\n",
    "batch_size = 100\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# 模型定义（修正输入维度）\n",
    "class RNNModel(nn.Module):\n",
    "    def __init__(self, input_dim=7781, hidden_dim=100, layer_dim=1, output_dim=10):\n",
    "        super().__init__()\n",
    "        self.rnn = nn.RNN(\n",
    "            input_size=input_dim,\n",
    "            hidden_size=hidden_dim,\n",
    "            num_layers=layer_dim,\n",
    "            batch_first=True\n",
    "        )\n",
    "        self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        h0 = torch.zeros(self.rnn.num_layers, batch_size, self.rnn.hidden_size).requires_grad_()\n",
    "        out, _ = self.rnn(x, h0.detach())\n",
    "        return self.fc(out[:, -1, :])\n",
    "\n",
    "# 初始化模型和优化器\n",
    "model = RNNModel(input_dim=7781)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# 训练循环\n",
    "num_epochs = 30\n",
    "best_accuracy = 0.0\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    for i, (inputs, labels) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    # 验证阶段\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in test_loader:\n",
    "            outputs = model(inputs)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    \n",
    "    accuracy = 100 * correct / total\n",
    "    print(f'Epoch {epoch+1:03d} | Val Acc: {accuracy:.2f}%')\n",
    "    \n",
    "    if accuracy > best_accuracy:\n",
    "        best_accuracy = accuracy\n",
    "        torch.save(model.state_dict(), 'best_model.pth')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
